import json
import os
import re
import platform
import cv2
import numpy as np
import torch
from torch.utils.data import Dataset
import Utils
from Utils import *
import torchvision.transforms.functional as TF
from torchvision import transforms
import random
import torch.nn.functional as F
import re
from Datasets.DatasetRegister import DefaultDatasetRegister as DatasetRegister

@DatasetRegister.register("FormatDataset")
class FormatDataset(Dataset):
    def __init__(self, dir_path, abs_flag=True, transform=None, random_pair = False, from_cache=True, update_cache=True):
        super().__init__()
        self.dir_path = dir_path
        self.abs_flag = abs_flag
        self.transform = transform if transform != None else transforms.ToTensor()
        self.from_cache = from_cache
        self.update_cache = update_cache
        self.random_pair = random_pair
        self.test_dataset_key = ""
        if "testRawData_incoherent_diffProtocol_binning_cut" in self.dir_path:
            self.test_dataset_key = "diff"
        elif "testRawData_incoherent_sameProtocol_binning_cut" in self.dir_path:
            self.test_dataset_key = "same"
        self.data_list = self.get_data_list(dir_path,self.from_cache,self.update_cache)
        random.shuffle(self.data_list)

    def get_key(self, path):
        if self.test_dataset_key in ["same", "diff"]:
            return get_parent_dirname(os.path.split(path)[0]) + "_" + os.path.basename(path)
        return ""

    def get_z_stack(self, path):
        parent_dir_path = os.path.split(path)[0]
        z_stack = os.listdir(parent_dir_path)
        z_stack = [os.path.join(parent_dir_path, name) for name in z_stack]
        z_stack = [path for path in z_stack if is_img(path)]
        z_stack.sort(key=lambda x: self.getPos(x))
        return z_stack

    def get_data_list(self, dir_path, from_cache, update_cache):
        cache_file_path = os.path.join(dir_path, "data_list.cache.json")
        if from_cache and os.path.exists(cache_file_path):
            data_list = json.load(open(cache_file_path, "r"))
            return data_list

        data_list = []
        for path in listAllImagePath(dir_path):
            pos_nm = self.getPos(path)
            label = self.getLabel(path)
            key = self.get_key(path)
            z_stack = self.get_z_stack(path)
            z_stack_index = z_stack.index(path)
            data_list.append({"img_path":path, "pos_nm":pos_nm,
                              "label":label, "key":key, "z_stack":z_stack,
                              "z_stack_index":z_stack_index,
                             })
        if update_cache:
            json.dump(data_list, open(cache_file_path, "w"), indent=4)
        return data_list

    def getPos(self, path):
        basename = os.path.basename(path)
        pos_nm = int(basename.split(".")[0])
        return pos_nm

    def getLabel(self, path):
        label = self.getPos(path)
        if self.abs_flag:
            label = abs(label)
        return label

    def __len__(self):
        return len(self.data_list)

    def __getitem__(self, index):
        data_item = self.data_list[index]
        img_path = data_item["img_path"]
        label = data_item["label"]
        key = data_item["key"]
        label = torch.Tensor([label])
        img = cv_imread(img_path)
        otherInfoDict = {"key":key, "test_dataset_key":self.test_dataset_key}
        if self.random_pair:
            z_stack = data_item["z_stack"]
            random_index = random.randint(0, len(z_stack) - 1)
            random_path = z_stack[random_index]
            random_img = cv_imread(random_path)
            # print("before same_transform_with_imgs",random_path, random_img.shape, img.shape)
            img, random_img = same_transform_with_imgs([img, random_img], self.transform)
            # print("after same_transform_with_imgs")
            return img, random_img, label, otherInfoDict
        img = self.transform(img)
        return img, label, otherInfoDict

if __name__ == '__main__':
    pass



